import io
from random import randbytes  # noqa: DUO102

from common import OperatingSystem
from common.utils.file_utils import append_bytes

from . import IAgentBinaryRepository

DEFAULT_NULL_BYTES_LENGTH = 16
DEFAULT_RANDOM_BYTES_LENGTH = 8


class PolymorphicAgentBinaryRepositoryDecorator(IAgentBinaryRepository):
    """
    PolymorphicAgentBinaryRepositoryDecorator adds random bytes to agent binaries in a way that does
    not affect their functionality. This allows Infection Monkey to emulate a property of
    polymorphic malware: all copies have different hashes.
    """

    def __init__(
        self,
        agent_binary_repository: IAgentBinaryRepository,
        null_bytes_length: int = DEFAULT_NULL_BYTES_LENGTH,
        random_bytes_length: int = DEFAULT_RANDOM_BYTES_LENGTH,
    ):
        self._agent_binary_repository = agent_binary_repository
        self._random_bytes_length = random_bytes_length
        self._null_bytes = b"\x00" * null_bytes_length

    def get_agent_binary(self, operating_system: OperatingSystem) -> io.BytesIO:
        agent_binary = self._agent_binary_repository.get_agent_binary(operating_system)

        # Note: These null bytes separate the Agent binary from the masque. The goal is to prevent
        # the masque from being interpreted by the OS as code that should be run.
        append_bytes(agent_binary, self._null_bytes + randbytes(self._random_bytes_length))

        return agent_binary
